from typing import Union, Dict, Tuple, List
import xlsxwriter

from io import BytesIO
import os
import string
import pandas as pd


class Style:
    """ 格式类 """

    def __init__(
            self,
            title_style: Union[Dict] = None,
            content_style: Union[Dict] = None,
            merge_style: dict = None,
            sheet_name: Union[str] = 'Sheet1'
    ):
        local_title_style = {
            'font': '仿宋',
            'align': 'center',
            'valign': 'center',
            'bold': True,
            'border': 1,
            'font_size': 12
        }
        self.title_style = local_title_style if not title_style else title_style
        local_content_style = {
            'font': '仿宋',
            'align': 'center',
            'valign': 'center'
        }
        self.content_style = local_content_style if not content_style else content_style
        self.merge_style = merge_style if merge_style else {
            'bold': True,
            'border': 1,
            'align': 'center',
            'valign': 'vcenter',
            'font': '仿宋',
        }
        self.sheet_name = sheet_name

        # merge和位置索引的关系
        self.letter = [item for item in string.ascii_uppercase]

    def cell_format(self, wb):
        return wb.add_format(self.title_style), wb.add_format(self.content_style)


class DTBase(Style):
    """
    行写入或者列写入数据
    """

    def __init__(
            self,
            columns: Union[Tuple, List],  # 表头
            content_w_row: Union[int] = 1,  # 预留表头位置 默认从第二行开始写入
            column_width: int = None,  # 设置列宽 默认不设置
            col_widths: List[int] = None,  # 列宽度列表
            is_w_col: Union[bool] = True,  # 是否列写入数据 # 默认列
            data: Union[Tuple, List] = None,  # 数据
            *args, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.content_w_row = content_w_row
        self.column_width = column_width
        self.col_widths = col_widths
        self.is_w_col = is_w_col

        self.data = data if data else []
        self.columns = columns

    def write_column(self, ws, data, content_format):
        """列写入"""
        for index, item in enumerate(data):
            ws.write_column(self.content_w_row, index, item, cell_format=content_format)

    def write_row(self, ws, data, content_format):
        """行写入"""
        for index, item in enumerate(data):
            ws.write_row(self.content_w_row + index, 0, item, cell_format=content_format)

    def set_column(self, ws, col_i):
        """
        设置列宽 每行设置列宽 或 统一列宽
        :param ws:
        :param col_i:
        :return:
        """
        # 统一设置列宽
        if self.column_width:
            ws.set_column(0, col_i, self.column_width)
        elif self.col_widths:
            for index in range(col_i):
                ws.set_column(index, index, self.col_widths[index])

    def worksheet(self, wb):
        ws = wb.add_worksheet(self.sheet_name)
        # 生成 样式
        columns_format, content_format = self.cell_format(wb)

        # 写入 头
        ws.write_row(0, 0, self.columns, cell_format=columns_format)
        # 写入 内容
        if self.is_w_col:
            self.write_column(ws, self.data, content_format)
        else:
            self.write_row(ws, self.data, content_format)

        if self.column_width:
            self.set_column(ws, len(self.columns) - 1)


class DTFile(DTBase):
    """生成文件"""

    def __init__(self,
                 file_name: Union[str],
                 path: Union[str],
                 *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.file_path = os.path.join(path, file_name)

    def write(self):
        with xlsxwriter.Workbook(self.file_path, {'in_memory': True}) as wb:
            self.worksheet(wb)


class DTBytesIO(DTBase):
    """ 文件流 """

    def write(self):
        output = BytesIO()
        with xlsxwriter.Workbook(output) as wb:
            self.worksheet(wb)
        return output.getvalue()


class PDBase(Style):
    """pandas to excel 修改样式 """

    def __init__(
            self,
            df: pd.DataFrame,
            columns: Union[List, Tuple] = None,
            col_widths: List[int] = None,  # 列宽度列表
            column_width: int = None,  # 列宽度列表
            *args,
            **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.df = df
        self.columns = columns
        self.col_widths = col_widths
        self.column_width = column_width

    def set_style(self, writer: pd.ExcelWriter, f=None, **kwargs):
        """设置样式"""
        workbook = writer.book
        worksheet = writer.sheets[self.sheet_name]
        # 生成样式
        columns_format, content_format = self.cell_format(workbook)

        columns = self.df.columns.values if not self.columns else self.columns
        for col_num, value in enumerate(columns):
            worksheet.write(0, col_num, value, columns_format)

        # 存在单独处理
        if f is not None:
            merge_format = workbook.add_format(self.merge_style)
            f(worksheet, merge_format, self, **kwargs)

        # 设置样式
        self.set_column(worksheet, len(self.df.columns) - 1)

    def set_column(self, ws, col_i):
        """
        设置列宽 每行设置列宽 或 统一列宽
        :param ws:
        :param col_i:
        :return:
        """
        # 统一设置列宽
        if self.column_width:
            ws.set_column(0, col_i, self.column_width)
        elif self.col_widths:
            for index in range(col_i):
                ws.set_column(index, index, self.col_widths[index])

    @staticmethod
    def merge_cell(worksheet, merge_format, self, **kwargs):
        """
        合并单元格  26 字母 不能超过
        :param worksheet:
        :param merge_format:
        :param self:
        :return:
        """
        col_num, row_num = len(self.columns), len(self.df)
        col_letter = self.letter[col_num - 1]

        worksheet.merge_range('A%d:%s%d' % (row_num, col_letter, row_num), kwargs['content'], merge_format)


class PDFile(PDBase):
    """ DataFrame 生成文件 """

    def __init__(
            self,
            path: str,
            *args, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.file_path = path

    def write(self, f: None, **kwargs):
        writer = pd.ExcelWriter(self.file_path, engine='xlsxwriter')
        self.df.to_excel(writer, index=False, header=False, sheet_name=self.sheet_name)
        self.set_style(writer, f, **kwargs)
        writer.save()


class PDBytes(PDBase):
    """ DataFrame 文件流 """

    def writer(self):
        output = BytesIO()
        writer = pd.ExcelWriter(output, engine='xlsxwriter')
        self.df.to_excel(writer, index=False, header=False, sheet_name=self.sheet_name)
        self.set_style(writer)
        writer.save()
        return output.getvalue()
